#include "limbo/errc.h"
#include "limbo/icmp/destination_unreachable.hpp"
#include "limbo/icmp/echo_request.hpp"
#include "limbo/icmp/echo_response.hpp"
#include "limbo/icmp/state.hpp"
#include "limbo/ip/state.hpp"
#include "limbo/stack.hpp"
#include "test-utils.h"

using namespace limbo;

using IP = ip::State<void, ip::Opts::calculate_checksum>;
using ICMP = icmp::State<IP>;
using MyStack = Stack<Layer<ICMP>, Layer<IP>>;

unsigned char example_raw[] = {
    0x45, 0x00, 0x00, 0x54, 0x77, 0x3a, 0x40, 0x00, 0x40, 0x01, 0xc5, 0x6c,
    0x7f, 0x00, 0x00, 0x01, 0x7f, 0x00, 0x00, 0x01, 0x08, 0x00, 0x00, 0x00,
    0xef, 0x13, 0x00, 0x01, 0x52, 0xa0, 0x4f, 0x63, 0x00, 0x00, 0x00, 0x00,
    0x9c, 0x82, 0x0c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x11, 0x12, 0x13,
    0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2a, 0x2b,
    0x2c, 0x2d, 0x2e, 0x2f, 0x30, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37,
};
auto example_chunk = Chunk(example_raw, sizeof(example_raw));
auto payload_chunk = Chunk((void *)(example_raw + 24), 60);

auto src = make_address("127.0.0.1");
auto dst = src;

TEST_CASE("send", "[icmp+ip]") {
  auto stack = MyStack();
  auto &ip_state = stack.get<1, 0>();
  auto &icmp_state = stack.get<0, 0>();
  ip_state.init(0x7739);

  auto ip_ctx = IP::Context{src, dst, 64, ip::Proto::icmp, 0b010, 0, nullptr};
  auto icmp_ctx = ICMP::Context{icmp::Type::echo_request, 0, &ip_ctx};
  unsigned char buff[sizeof(example_raw)];

  SECTION("no enough space") {
    auto zero_chunk = Chunk(buff, 0);
    auto result = stack.send(zero_chunk, icmp_ctx, payload_chunk);
    REQUIRE(!result);
    CHECK(result.state() == &icmp_state);
    CHECK(result.error_code() == (uint32_t)Errc::no_enough_space);
  }

  auto buff_chunk = Chunk(buff, sizeof(buff));
  auto result = stack.send(buff_chunk, icmp_ctx, payload_chunk);

  REQUIRE(result);
  auto res_buff = result.consumed();
  CHECK((void *)res_buff.data() == (void *)buff);
  CHECK(res_buff.size() == example_chunk.size());
  CHECK(res_buff == example_chunk);
}

TEST_CASE("receive", "[icmp+ip]") {
  auto stack = MyStack();
  auto &icmp_state = stack.get<0, 0>();
  auto result = stack.recv(example_chunk, nullptr);
  REQUIRE(result);
  REQUIRE(result.consumed() == example_chunk);
  REQUIRE(result.state() == &icmp_state);
  auto &p = icmp_state.get_parsed();
  CHECK(p.type == icmp::Type::echo_request);
  CHECK(p.code == 0);
  CHECK(p.payload == payload_chunk);
  CHECK(p.container->source == src);
  CHECK(p.container->destination == dst);
  CHECK(p.container->flag == 0b010);
}

TEST_CASE("receive, checksums manipulations", "[icmp+ip]") {
  using ICMP = icmp::State<IP, icmp::Opts::validate_checksum>;
  using MyStack = Stack<Layer<ICMP>, Layer<IP>>;
  unsigned char buff[sizeof(example_raw)];
  std::memcpy(buff, example_raw, example_chunk.size());
  auto chunk = Chunk(buff, sizeof(buff));

  SECTION("validate correct checksum") {
    auto stack = MyStack();
    auto &icmp_state = stack.get<0, 0>();
    buff[22] = 0xFF;
    buff[23] = 0x91;

    auto result = stack.recv(chunk, nullptr);
    REQUIRE(result);
    REQUIRE(result.consumed() == chunk);
    REQUIRE(result.state() == &icmp_state);

    auto &p = icmp_state.get_parsed();
    CHECK(p.type == icmp::Type::echo_request);
    CHECK(p.code == 0);
    CHECK(p.payload == payload_chunk);
    CHECK(p.container->source == src);
    CHECK(p.container->destination == dst);
    CHECK(p.container->flag == 0b010);
  }

  SECTION("invalid checksum can be ignored") {
    using ICMP = icmp::State<IP>;
    using MyStack = Stack<Layer<ICMP>, Layer<IP>>;
    auto stack = MyStack();
    auto &icmp_state = stack.get<0, 0>();
    auto orig = buff[0x24];
    buff[0x24] = 11;
    auto result = stack.recv(chunk, nullptr);
    buff[0x24] = orig;

    REQUIRE(result);
    REQUIRE(result.consumed() == example_chunk);
    REQUIRE(result.state() == &icmp_state);

    auto &p = icmp_state.get_parsed();
    CHECK(p.type == icmp::Type::echo_request);
    CHECK(p.code == 0);
    CHECK(p.payload == payload_chunk);
    CHECK(p.container->source == src);
    CHECK(p.container->destination == dst);
    CHECK(p.container->flag == 0b010);
  }
}

TEST_CASE("receive, wrong checksum", "[icmp+ip]") {
  using ICMP = icmp::State<IP, icmp::Opts::validate_checksum>;
  using MyStack = Stack<Layer<ICMP>, Layer<IP>>;

  unsigned char buff[sizeof(example_raw)];
  std::memcpy(buff, example_raw, example_chunk.size());
  buff[0x24] = 0x12;

  auto chunk = Chunk(buff, sizeof(buff));
  auto stack = MyStack();
  auto &icmp_state = stack.get<0, 0>();

  auto result = stack.recv(chunk, nullptr);
  REQUIRE(!result);
  REQUIRE(result.state() == &icmp_state);
  CHECK(result.error_code() == (uint32_t)Errc::icmp_checksum_mismatch);
}

TEST_CASE("echo req/res", "[icmp+ip]") {
  using IP = ip::State<void>;
  using ICMP = icmp::State<IP>;
  using EchoReq = icmp::EchoRequest<ICMP>;
  using EchoRes = icmp::EchoResponse<ICMP>;
  using MyStack = Stack<Layer<EchoRes, EchoReq>, Layer<ICMP>, Layer<IP>>;

  MyStack stack;
  auto &echo_req = stack.get<0, 1>();
  auto &echo_res = stack.get<0, 0>();
  echo_req.init(src);
  echo_res.init(src);
  unsigned char payload_raw[4] = {1, 2, 3, 4};
  auto payload = Chunk(payload_raw, sizeof(payload_raw));
  char tx_buff_raw[32];
  char rx_buff_raw[32];
  auto tx_buff = Chunk(tx_buff_raw, sizeof(tx_buff_raw));
  auto rx_buff = Chunk(rx_buff_raw, sizeof(rx_buff_raw));

  auto tx_result = echo_req.send(stack, tx_buff, dst, payload);
  REQUIRE(tx_result);
  CHECK(tx_result.consumed() == tx_buff);

  SECTION("recv req") {
    auto result = stack.recv(tx_buff, nullptr);
    REQUIRE(result);
    REQUIRE(!result.demand());
    REQUIRE(result.state() == &echo_req);
    auto &p = echo_req.get_parsed();
    CHECK(p.identifier == 1);
    CHECK(p.sequence_no == 1);
    CHECK(p.payload == payload);

    CHECK(p.container->type == icmp::Type::echo_request);
    CHECK(p.container->code == 0);
    CHECK(p.container->container->source == src);
    CHECK(p.container->container->destination == dst);
    CHECK(p.container->container->flag == 0);
  }

  auto rx_result = echo_res.send(stack, rx_buff, echo_req.get_parsed());
  REQUIRE(rx_result);
  CHECK(rx_result.consumed() == rx_buff);

  auto result = stack.recv(rx_buff, nullptr);
  REQUIRE(result);
  REQUIRE(!result.demand());
  REQUIRE(result.state() == &echo_res);
  auto &p = echo_res.get_parsed();
  CHECK(p.identifier == 1);
  CHECK(p.sequence_no == 1);
  CHECK(p.payload == payload);

  CHECK(p.container->type == icmp::Type::echo_response);
  CHECK(p.container->code == 0);
  CHECK(p.container->container->source == src);
  CHECK(p.container->container->destination == dst);
  CHECK(p.container->container->flag == 0);
}

TEST_CASE("destination unreacheable req/res", "[icmp+ip]") {
  using IP = ip::State<void, ip::Opts::calculate_checksum>;
  using ICMP = icmp::State<IP, icmp::Opts::calculate_checksum>;
  using Unreach = icmp::DestinationUnreacheable<ICMP>;
  using MyStack = Stack<Layer<Unreach>, Layer<ICMP>, Layer<IP>>;

  uint8_t ip_data[] = {
      0x45, 0x00, 0x00, 0x1f, 0x57, 0x52, 0x40, 0x00, 0x40, 0x11, 0x9a,
      0x1a, 0xc0, 0xa8, 0x64, 0x0f, 0xc0, 0xa8, 0x64, 0x01, 0x19, 0xae,
      0x04, 0xd2, 0x00, 0x0b, 0xa3, 0x7b, 0x7a, 0x7a, 0x7a,
  };
  auto ip_chunk = Chunk(ip_data, sizeof(ip_data));

  uint8_t example_data[] = {
      0x45, 0xc0, 0x00, 0x3b, 0x10, 0x6e, 0x00, 0x00, 0x40, 0x01, 0x20, 0x33,
      0xc0, 0xa8, 0x64, 0x01, 0xc0, 0xa8, 0x64, 0x0f, 0x03, 0x03, 0x46, 0x7b,
      0x00, 0x00, 0x00, 0x00, 0x45, 0x00, 0x00, 0x1f, 0x57, 0x52, 0x40, 0x00,
      0x40, 0x11, 0x9a, 0x1a, 0xc0, 0xa8, 0x64, 0x0f, 0xc0, 0xa8, 0x64, 0x01,
      0x19, 0xae, 0x04, 0xd2, 0x00, 0x0b, 0xa3, 0x7b, 0x7a, 0x7a, 0x7a,
  };
  auto example_chunk = Chunk(example_data, sizeof(example_data));

  MyStack stack;
  auto src = make_address("192.168.100.1");
  auto dst = make_address("192.168.100.15");
  auto &ip = stack.get<2, 0>();
  auto &unreach = stack.get<0, 0>();

  ip.init(0x106d);
  unreach.init(src, 0xc0);

  auto r = stack.recv(ip_chunk, nullptr);
  REQUIRE(r);
  REQUIRE(r.state() == &ip);
  REQUIRE(r.consumed() == ip_chunk);

  uint8_t buff_raw[59];
  auto buff = Chunk(buff_raw, sizeof(buff_raw));

  auto &ip_packet = ip.get_parsed();
  r = unreach.send(stack, buff, ip_packet, 3);
  REQUIRE(r);
  CHECK(r.consumed().size() == buff.size());
  CHECK(r.consumed() == example_chunk);

  unreach.init(dst);
  r = stack.recv(example_chunk, nullptr);
  REQUIRE(r);
  REQUIRE(!r.demand());
  REQUIRE(r.state() == &unreach);

  auto &icmp_packet = unreach.get_parsed();
  CHECK(icmp_packet.type == icmp::Type::destination_unreachable);
  CHECK(icmp_packet.code == 3);
  auto payload = Chunk(example_data + 24, example_chunk.size() - 24);
  CHECK(icmp_packet.payload == payload);
  CHECK(icmp_packet.container->source == src);
  CHECK(icmp_packet.container->destination == dst);
}
